import random

import torchvision
from skimage import io,transform
import os
import pickle
import torch
import numpy as np
import torch.utils.data as data
import matplotlib.pyplot as plt
from PIL import Image
from torchvision import datasets, transforms


class MyDatasets(data.Dataset):#继承data.Dataset
    #初始化函数
    def __init__(self,Input_root_dir,Target_root_dir,start_point_y=0,data_size=[256,128],transform=None):

        #输入数据文件目录
        self.Input_root_dir = Input_root_dir
        # 目标数据文件目录
        self.Target_root_dir = Target_root_dir
        #变换
        self.transform = transform
        #目录里的所有文件
        self.Inputdata = os.listdir(self.Input_root_dir)
        self.Targetdata = os.listdir(self.Target_root_dir)

        self.start_point_y = start_point_y
        self.Data_size = data_size

    #返回整个输入数据集的大小
    def __len__(self):
        return len(self.Targetdata)

    #根据索引返回对应文件
    def __getitem__(self, index):
        #根据索引item获取该输入数据
        data_index = self.Targetdata[index]

        # 确定矩阵初始截取y点
        start_point_y = self.start_point_y

        #确定训练数据大小
        data_size = self.Data_size


        #返回数据路径
        # data_path = os.path.join(self.Input_root_dir,data_index)

        #读取数据
        # data_input = (np.load(data_path,allow_pickle=True))
        # (l,w) = data_input.shape
        # data_input = data_input.reshape((l,w,1))
        # print(data_input.shape)


        # data_input = data_input[start_point_y:start_point_y+data_size[0],start_point_y:start_point_y+data_size[1]]

        # noiseLevel , dataIndex = data_index.split('.')[0].split('_')[-1].split('NoiseData')
        # data_target_name = 'theoreticalData'+str(noiseLevel)+'_'+str(dataIndex)+'.npy'

        data_target = (np.load(os.path.join(self.Target_root_dir,data_index),allow_pickle=True))
        SNR = (random.sample([-5,-1,1,5,10], 1))[0]
        noise = np.random.randn(data_target.shape[0],data_target.shape[1]) 	#产生N(0,1)噪声数据
        noise = noise - np.mean(noise) 								#均值为0
        signal_power = np.linalg.norm(data_target - data_target.mean() )**2 / data_target.size	#此处是信号的std**2
        noise_variance = signal_power/np.power(10,(SNR/10))         #此处是噪声的std**2
        noise = (np.sqrt(noise_variance) / np.std(noise) ) * noise    ##此处是噪声的std**2
        data_input = noise + data_target
        (l,w) = data_input.shape
        data_input = data_input.reshape((l,w,1))
        # data_target = data_target[start_point_y:start_point_y+data_size[0],start_point_y:start_point_y+data_size[1]]
        data_target = noise
        (l,w) = data_target.shape
        data_target = data_target.reshape((l,w,1))

        if self.transform:
            data_input = self.transform(data_input)
            data_target = self.transform(data_target)

        # print(data_input)
        # data_input.unsqueeze(0)
        # data_target.unsqueeze(0)
        # print(data_input.shape)

        return data_input,data_target

if __name__ == '__main__':


    Input_root_dir = r'Data\Numpy_DATA\AddNoise'
    Target_root_dir =r'Data\Numpy_DATA\Original'
    data = MyDatasets(Input_root_dir=Input_root_dir,Target_root_dir=Target_root_dir,transform=None)
    dataloader = torch.utils.data.DataLoader(data,batch_size=1,shuffle=True)

    input_example, target_example,data_target_name,noiseLevel = next(iter(dataloader))
    print(u'X_example:{}'.format((input_example)))
    print(u'Y_example:{}'.format((target_example)))
    print(data_target_name)
    print(noiseLevel)



